package com.ht.modules.component.ai.assistants;

import com.alibaba.fastjson2.JSON;
import com.ht.modules.component.ai.mcp.entity.McpServers;
import com.ht.modules.component.ai.mcp.entity.McpServersInfo;
import com.ht.modules.component.ai.mcp.handler.OwnerMcpLogMessageHandler;
import com.ht.modules.component.ai.session.service.SessionChatStoreService;
import com.ht.modules.component.ai.tools.service.ToolService;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.mcp.McpToolProvider;
import dev.langchain4j.mcp.client.DefaultMcpClient;
import dev.langchain4j.mcp.client.McpClient;
import dev.langchain4j.mcp.client.transport.McpTransport;
import dev.langchain4j.mcp.client.transport.http.HttpMcpTransport;
import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.StreamingChatModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.service.tool.ToolProvider;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.pgvector.PgVectorEmbeddingStore;
import org.apache.commons.io.IOUtils;
import org.apache.logging.log4j.util.Strings;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.io.ClassPathResource;

import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

/**
 * 助手服务
 *
 * @since 2024/1/17
 */
@Configuration
public class AssistantService<T> implements IAssistantService<T> {
    @Value("${spring.data.pgvector.host}")
    private String host;
    @Value("${spring.data.pgvector.port}")
    private int port;
    @Value("${spring.data.pgvector.database}")
    private String database;
    @Value("${spring.data.pgvector.user}")
    private String user;
    @Value("${spring.data.pgvector.password}")
    private String password;
    @Value("${spring.data.pgvector.createTable}")
    private Boolean createTable;
    @Value("${spring.data.pgvector.dropTableFirst}")
    private Boolean dropTableFirst;
    @Autowired
    private EmbeddingModel embeddingModel;
    @Autowired
    private SessionChatStoreService sessionChatStoreService;

    /**
     * 构建助手
     *
     * @param streamingChatModel 流式聊天模型
     * @param chatMemoryProvider 聊天记忆
     * @param contentRetriever   内容检索器
     * @param toolService        工具服务
     * @param clazz              助手类
     * @return T 助手
     */
    @Override
    public T buildAssistant(StreamingChatModel streamingChatModel,
                            ChatMemoryProvider chatMemoryProvider,
                            ContentRetriever contentRetriever,
                            ToolService toolService,
                            ToolProvider toolProvider,
                            Class<T> clazz
    ) {
        return AiServices.builder(clazz)
                .streamingChatModel(streamingChatModel)
                .chatMemoryProvider(chatMemoryProvider)
                .contentRetriever(contentRetriever)
                .tools(toolService)
                .toolProvider(toolProvider)
                .build();
    }

    /**
     * 构建工具提供器
     *
     * @param resourceLocation 资源文件位置
     * @return ToolProvider 工具提供器
     */
    public ToolProvider buildToolProvider(String resourceLocation) {
        try {
            ClassPathResource classPathResource = new ClassPathResource(resourceLocation);
            String mcpServersJson = IOUtils.toString(classPathResource.getInputStream(), StandardCharsets.UTF_8);
            if (Strings.isBlank(mcpServersJson)) {
                throw new RuntimeException(resourceLocation + " is empty");
            }
            McpServers mcpServers = JSON.parseObject(mcpServersJson, McpServers.class);
            if (mcpServers == null) {
                throw new RuntimeException(resourceLocation + " ParseObject is empty");
            }
            Map<String, McpServersInfo> mcpServersInfoMap = mcpServers.getMcpServers();
            if (mcpServersInfoMap == null) {
                throw new RuntimeException(resourceLocation + " McpServersInfoMap is empty");
            }
            List<McpClient> mcpClients = new ArrayList<>();
            mcpServersInfoMap.forEach((name, mcpServersInfo) -> {
                if (!mcpServersInfo.getIsActive()) {
                    return;
                }
                McpTransport mcpTransport = new HttpMcpTransport.Builder()
                        .sseUrl(mcpServersInfo.getUrl())
                        .timeout(Duration.ofSeconds(600))
                        .logRequests(mcpServersInfo.getLogRequests())
                        .logResponses(mcpServersInfo.getLogResponses())
                        .build();
                McpClient mcpClient = new DefaultMcpClient.Builder()
                        .transport(mcpTransport)
                        .resourcesTimeout(Duration.ofSeconds(600))
                        .toolExecutionTimeoutErrorMessage("工具执行超时，请稍后再试")
                        .initializationTimeout(Duration.ofSeconds(600))
                        .promptsTimeout(Duration.ofSeconds(600))
                        .pingTimeout(Duration.ofSeconds(600))
                        .toolExecutionTimeout(Duration.ofSeconds(600))
                        .logHandler(new OwnerMcpLogMessageHandler())
                        .build();
                mcpClients.add(mcpClient);
            });
            return McpToolProvider.builder()
                    .mcpClients(mcpClients)
                    .build();
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 构建向量数据库内容检索器
     *
     * @param tableName         内容检索表名
     * @param maxResults        最大结果数
     * @param minScore          最小分数
     * @param dynamicMinScore   动态最小分数
     * @param dynamicMaxResults 动态最大结果数
     * @return EmbeddingStoreContentRetriever 向量数据库内容检索器
     */
    @Override
    public EmbeddingStoreContentRetriever buildEmbeddingStoreContentRetriever(String tableName,
                                                                              int maxResults,
                                                                              double minScore,
                                                                              double dynamicMinScore,
                                                                              int dynamicMaxResults
    ) {
        return EmbeddingStoreContentRetriever.builder()
                .embeddingStore(this.buildEmbeddingStore(tableName))
                .embeddingModel(this.embeddingModel)
                .maxResults(maxResults)
                .dynamicMaxResults(query -> dynamicMaxResults)
                .minScore(minScore)
                .dynamicMinScore(query -> dynamicMinScore)
                .build();
    }


    /**
     * 构建向量数据库
     *
     * @param tableName 表名
     * @return EmbeddingStore<TextSegment> 向量数据库
     */
    @Override
    public EmbeddingStore<TextSegment> buildEmbeddingStore(String tableName) {
        return PgVectorEmbeddingStore.builder()
                .host(host)
                .port(port)
                .database(database)
                .user(user)
                .password(password)
                .table(tableName)
                .dimension(embeddingModel.dimension())
                .createTable(createTable)
                .dropTableFirst(dropTableFirst)
                .build();
    }

    /**
     * 构建聊天记忆
     *
     * @param maxMessages 最大消息数
     * @return ChatMemoryProvider 聊天记忆
     */
    @Override
    public ChatMemoryProvider buildChatMemoryProvider(int maxMessages) {
        return (memoryId) -> MessageWindowChatMemory.builder()
                .id(memoryId)
                .maxMessages(maxMessages)
                .chatMemoryStore(sessionChatStoreService)
                .build();
    }
}
